package avltree;

import javafx.util.Pair;

public class AVLTree {
    static class TreeNode {
        public int val;
        public int bf;//平衡因子
        TreeNode left;//左子树
        TreeNode right;//右子树
        TreeNode parent;//父亲结点

        public TreeNode(int val) {
            this.val = val;
        }
    }
    public TreeNode root;//根节点

    public boolean insert(int val) {
        TreeNode node = new TreeNode(val);
        if(root == null) {
            root = node;
            return true;
        }
        TreeNode parent = null;
        TreeNode cur = root;
        while(cur != null) {
            if(cur.val > val) {
                //向左寻找
                parent = cur;
                cur = cur.left;
            } else if(cur.val == val) {
                //相等就说明插入失败
                return false;
            } else {
                //向右寻找
                parent = cur;
                cur = cur.right;
            }
        }
        //cur == null
        if(parent.val > val) {
            parent.left = node;
        } else {
            parent.right = node;
        }
        node.parent = parent;
        cur = node;
        //调整因子
        while(parent != null) {
            //新增结点在右树，平衡因子++，新增结点在左树, 平衡因子--
            if(cur == parent.right) {
                //新增结点在右树
                parent.bf++;
            } else {
                //新增结点在左树
                parent.bf--;
            }

            //检查平衡因子大小
            if(parent.bf == 0) {
                //这里说明以及平衡了，不需要在调整
                break;
            } else if(parent.bf == 1 || parent.bf == -1) {
                //再往上查看结点，有可能不平衡
                cur = parent;
                parent = cur.parent;
            } else {
                //一定不平衡
                if(parent.bf == 2) {//右树高，需要降低右树的高度
                    if(cur.bf == 1) {
                        //左旋
                        rotateLeft(parent);
                    } else {//cur.bf == -1
                        //右左双旋
                        rottateRL(parent);
                    }
                } else {//parent.bf == -2  左树高，需要降低左树的高度
                    if(cur.bf == -1) {
                        //右旋
                        rotateRight(parent);
                    } else {//cur.bf == 1
                        //左右双旋
                        rortateLR(parent);
                    }
                }
                //走完以上代码即哦平衡
                break;
            }
        }
        return true;
    }

    //右左双旋
    private void rottateRL(TreeNode parent) {
        TreeNode subR = parent.right;
        TreeNode subRL = subR.left;
        int bf = subRL.bf;
        rotateRight(parent.right);
        rotateLeft(parent);
        if(bf == 1) {
            parent.bf = -1;
            subR.bf = 0;
            subRL.bf = 0;
        } else if(bf == -1) {
            subRL.bf = 0;
            parent.bf = 0;
            subR.bf = 1;
        }
    }

    //左右双旋
    private void rortateLR(TreeNode parent) {
        TreeNode subL = parent.left;
        TreeNode subLR = subL.right;
        int bf = subLR.bf;
        //选左旋后右旋
        rotateLeft(parent.left);//注意旋转的是parent的左子树
        rotateRight(parent);
        if(bf == -1) {
            subL.bf = 0;
            subLR.bf = 0;
            parent.bf = 1;
        } else if(bf == 1){
            subL.bf = -1;
            subLR.bf = 0;
            parent.bf= 0;
        } //bf == 0，这里本身就是平衡的，就不用修改
    }

    //左单旋
    private void rotateLeft(TreeNode parent) {
        TreeNode subR = parent.right;
        TreeNode subRL = subR.left;
        parent.right = subRL;

        subR.left = parent;

        if(subRL != null) {
            subRL.parent = parent;
        }
        //必须先记录parent的父亲
        TreeNode pParent = parent.parent;
        parent.parent = subR;
        //检查parent是否为根节点
        if(parent == root) {
            root = subR;
            root.parent = null;
        } else {
            //不是根节点需要判断parent是左子树还是右子树
            if(pParent.left == parent) {
                pParent.left = subR;
            } else {
                pParent.right = subR;
            }
            subR.parent = pParent;
        }
        subR.bf = 0;
        parent.bf = 0;
    }

    //右单旋
    private void rotateRight(TreeNode parent) {
        TreeNode subL = parent.left;
        TreeNode subLR = subL.right;
        parent.left = subLR;
        subL.right = parent;
        if(subLR != null) {
            subLR.parent = parent;
        }
        //必须先记录parent的父亲
        TreeNode pParent = parent.parent;
        parent.parent = subL;
        //检查parent是否是根节点
        if(parent == root) {
            root = subL;
            root = null;
        } else {
            //不是根节点，判断parent是左子树还是右子树
            if(pParent.left == parent) {
                pParent.left = subL;
            } else {
                pParent.right = subL;
            }
            subL.parent = pParent;
        }
        subL.bf = 0;
        parent.bf = 0;
    }
    //验证平衡
    public boolean isBalanced(TreeNode root) {
        if(root == null) {
            return true;
        }
        int leftH = heigh(root.left);
        int rightH = heigh(root.right);
        if(rightH - leftH != root.bf) {
            System.out.println("结点：" + root.val + "的平衡因子异常");
            return false;
        }
        return Math.abs(leftH - rightH) <= 1 &&
                isBalanced(root.left) &&
                isBalanced(root.right);
    }
    //求树的高度
    private int heigh(TreeNode root) {
        if(root == null) {
            return 0;
        }
        int leftH = heigh(root.left);
        int rightH = heigh(root.right);
        return leftH > rightH ? leftH + 1 : rightH + 1;
    }
    //验证顺序
    public void dfs(TreeNode root) {
        if(root == null) {
            return;
        }
        dfs(root.left);
        System.out.print("root = " + root.val + "; ");
        dfs(root.right);
    }
}
